use ncnnrs::*;

fn main() -> anyhow::Result<()> {
    println!("=== NCNN Rust Demo ===");

    // 1. 基本版本信息
    println!("\n1. 基本信息:");
    println!("NCNN版本: {}", version());

    // 2. GPU信息（如果支持）
    let gpu_count = get_gpu_count();
    println!("GPU数量: {}", gpu_count);

    if gpu_count > 0 {
        for i in 0..gpu_count {
            let device_name = get_device_name(i);
            let heap_budget = get_gpu_heap_budget(i);
            println!(
                "GPU {} - 设备名: {}, 堆预算: {} MB",
                i,
                device_name,
                heap_budget / 1024 / 1024
            );
        }
    }

    // 3. 矩阵操作演示
    println!("\n2. 矩阵操作演示:");
    demo_mat_operations()?;

    // 4. 网络操作演示
    println!("\n3. 网络操作演示:");
    demo_net_operations()?;

    // 5. 图像处理演示
    println!("\n4. 图像处理演示:");
    demo_image_processing()?;

    println!("\n=== Demo 完成 ===");
    Ok(())
}

fn demo_mat_operations() -> anyhow::Result<()> {
    // 创建不同维度的矩阵
    println!("  创建1D矩阵 (长度100):");
    let mat1d = Mat::new_1d(100, None);
    println!("  1D矩阵 - 维度: {}, 宽度: {}", mat1d.dims(), mat1d.w());

    println!("  创建2D矩阵 (32x32):");
    let mat2d = Mat::new_2d(32, 32, None);
    println!(
        "  2D矩阵 - 维度: {}, 宽度: {}, 高度: {}",
        mat2d.dims(),
        mat2d.w(),
        mat2d.h()
    );

    println!("  创建3D矩阵 (224x224x3, 常用于图像):");
    let mut mat3d = Mat::new_3d(224, 224, 3, None);
    println!(
        "  3D矩阵 - 维度: {}, 宽度: {}, 高度: {}, 通道: {}",
        mat3d.dims(),
        mat3d.w(),
        mat3d.h(),
        mat3d.c()
    );

    // 填充矩阵
    println!("  用值0.5填充3D矩阵");
    mat3d.fill(0.5);

    println!("  矩阵元素大小: {} bytes", mat3d.elemsize());
    println!("  矩阵元素包装: {}", mat3d.elempack());
    println!("  矩阵通道步长: {}", mat3d.cstep());

    Ok(())
}

fn demo_net_operations() -> anyhow::Result<()> {
    // 创建网络
    println!("  创建NCNN网络:");
    let mut net = Net::new();

    // 设置选项
    println!("  配置网络选项:");
    let mut opt = Option::new();
    opt.set_num_threads(4);
    opt.use_vulkan_compute(false); // 使用CPU模式
    net.set_option(&opt);

    println!("  网络创建成功! (注意: 这只是创建了空网络，没有加载模型)");
    println!("  要加载模型，可以使用:");
    println!("    net.load_param(\"model.param\")?;");
    println!("    net.load_model(\"model.bin\")?;");

    // 创建提取器
    println!("  创建网络提取器:");
    let extractor = net.create_extractor();
    println!("  提取器创建成功!");

    Ok(())
}

fn demo_image_processing() -> anyhow::Result<()> {
    // 创建模拟图像数据 (3通道RGB, 64x64像素)
    let width = 64i32;
    let height = 64i32;
    let channels = 3;
    let pixel_count = (width * height * channels) as usize;

    println!("  创建模拟RGB图像数据 ({}x{}, 3通道):", width, height);
    let mut image_data = vec![128u8; pixel_count]; // 用灰色填充

    // 创建一些模式
    for y in 0..height {
        for x in 0..width {
            let idx = ((y * width + x) * channels) as usize;
            image_data[idx] = (x * 255 / width) as u8; // R通道渐变
            image_data[idx + 1] = (y * 255 / height) as u8; // G通道渐变
            image_data[idx + 2] = 128; // B通道固定
        }
    }

    println!("  从像素数据创建Mat:");
    let mat = Mat::from_pixels(&image_data, MatPixelType::RGB, width, height, None)?;
    println!(
        "  图像Mat - 维度: {}, 宽度: {}, 高度: {}, 通道: {}",
        mat.dims(),
        mat.w(),
        mat.h(),
        mat.c()
    );

    println!("  创建缩放后的图像 (缩放到32x32):");
    let resized_mat =
        Mat::from_pixels_resize(&image_data, MatPixelType::RGB, width, height, 32, 32, None)?;
    println!(
        "  缩放后Mat - 宽度: {}, 高度: {}",
        resized_mat.w(),
        resized_mat.h()
    );

    // 演示图像预处理
    println!("  创建用于神经网络的图像Mat:");
    let mut nn_mat = Mat::from_pixels_resize(
        &image_data,
        MatPixelType::RGBtoBGR, // 转换为BGR（很多模型需要BGR格式）
        width,
        height,
        224,
        224, // 缩放到标准的224x224
        None,
    )?;

    // 标准化（减去均值并除以标准差）
    let mean_vals = [103.94f32, 116.78f32, 123.68f32]; // ImageNet均值
    let norm_vals = [0.017f32, 0.017f32, 0.017f32]; // 1/58.8 标准化
    nn_mat.substract_mean_normalize(&mean_vals, &norm_vals);

    println!(
        "  预处理完成 - 宽度: {}, 高度: {}, 通道: {}",
        nn_mat.w(),
        nn_mat.h(),
        nn_mat.c()
    );

    // 演示边界填充
    println!("  演示边界填充:");
    let padded_mat = copy_make_border(&nn_mat, 10, 10, 10, 10, BorderType::Constant, 0.0)?;
    println!(
        "  填充后Mat - 宽度: {}, 高度: {}",
        padded_mat.w(),
        padded_mat.h()
    );

    Ok(())
}
